#coding:UTF-8
'''
JAKA SDK Python API custom-made software development (CSD)
Date:
    2021-10-24 (Y-M-D)
Author:
    Xiaofeng Du, Shenghan Xie
Discription:
    JAKA robot movement control library.
'''

import os
import sys
import time
import numpy as np
# sys.path.append("add path of JAKA SDK")
sys.path.append("/home/yons/grasp_ws/JAKA_SDK/python3-refresh")
# import jaka sdk
try:
    import jkrc
except:
    raise NameError("JAKA SDK path error! current work path: "+os.getcwd())

class JAKA():
  # parameters
    _ABS = 0
    _INRC = 1
    _homePose = None
  # functions
    ''' Funtions list:
    # get pose
            # get [X, Y, Z] pose
        def getposXYZ(self)
            # get [Roll, Pitch, Yaw] pose
        def getposRPY(self)
            # get [X, Y, Z, Roll, Pitch, Yaw] pose
        def getpos6DoF(self)
    # connect & disconnect from the robot
            # connect
        def robot_disconnect(self)
            # disconnect
        def jaka_connect(self)
    # move
            # as name
        def moveInWorldCoordinate(self, INCRPose, speedInput = 50)
            # as name
        def moveInFlangeCoordinate(self, tarPose)
            # as name
        def moveInToolCoordinate(self, toolCorPose, tarPose)'''
    def __init__(self, address,connect = True):
        self.address = address
        self.robot = None
        if(connect):
            self.jaka_connect()

    def joint_move(self, joints,sp):
        #joints = 180*joints/pi
        self.robot.joint_move(joint_pos =joints, move_mode = 0,is_block= False ,speed = sp)
        time.sleep(0.08)

    def getjoints(self):
        ret = self.robot.get_joint_position()
        if ret[0] == 0:
            return ret[1]

    def liner_move(self, pos,sp):
        ret = self.robot.linear_move(end_pos = pos, move_mode = self._ABS, is_block = True, speed = sp)
        if(ret[0] == 0):
            pass
        else:
            print(ret[0])
        # time.sleep(0.08)
    # '''
    # INCRPose : increase position [X, Y, Z, Roll, Pitch, Yaw]
    # '''
    # def moveInWorldCoordinate(self, INCRPose, speedInput = 5):
    #     curPose = self.getpos6DoF()
    #     tarPose = [0,0,0,0,0,0]
    #     for i in range(3):
    #       INCRPose[i] = curPose[i]+INCRPose[i]
    #     ret = self.robot.linear_move(end_pos = INCRPose, move_mode = self._ABS, is_block = False, speed = SPEED)
    #     if(ret[0] == -4):
    #         print("Inverse solution failed")
    #     return ret

    # '''
    # tarPose : target position [X, Y, Z, Roll, Pitch, Yaw]
    # '''
    # def moveInFlangeCoordinate(self, tarPose):
    #     RPY = self.getposRPY()
    #     # get rotationMatrix from world to flange
    #     rotationMatrix = np.array(self.robot.rpy_to_rot_matrix(RPY)[1]) # w - end

    #     tarXYZ = np.array(tarPose[:3])
    #     tarRPY =np.array(tarPose[3:])

    #     tarRM = self.robot.rpy_to_rot_matrix(tarRPY)[1] # end - P

    #     worldXYZ = np.matmul(rotationMatrix,tarXYZ)

    #     worldRM = np.matmul(rotationMatrix,tarRM)
    #     print(worldRM)
    #     worldRPY = self.robot.rot_matrix_to_rpy(worldRM)[1]
    #     worldPose = np.hstack((worldXYZ,worldRPY))

    #     self.moveInWorldCoordinate(worldPose)

    # '''
    # toolCor : tool coordinate's position in flange coordinate  [X, Y, Z, Roll, Pitch, Yaw]
    # tarPose : target position in tool coordinate [X, Y, Z, Roll, Pitch, Yaw]
    # '''
    # def moveInToolCoordinate(self, toolCorPose, tarPose):
    #     # get rotationMatrix from tool to flange
    #     rotationMatrix = np.array(self.robot.rpy_to_rot_matrix(toolCorPose[3:])[1])
    #     tarXYZ = np.array(tarPose[:3])
    #     tarRPY =np.array(tarPose[3:])
    #     tarRM = self.robot.rpy_to_rot_matrix(tarRPY)[1]
    #     FlangeRM = np.matmul(rotationMatrix,tarRM)
    #     FlangeRPY = self.robot.rot_matrix_to_rpy(FlangeRM)[1]
    #     FlangeXYZ = np.matmul(rotationMatrix,tarXYZ) + toolCorPose[:3]
    #     FlangePose = np.hstack((FlangeXYZ,FlangeRPY))
    #     self.moveInFlangeCoordinate(FlangePose)

# get pose
    # get [X, Y, Z] pose
    def getposXYZ(self):
        ret = self.robot.get_tcp_position()
        if ret[0] == 0:
            return ret[1][0:3]
      # get [Roll, Pitch, Yaw] pose
    def getposRPY(self):
        ret = self.robot.get_tcp_position()
        if ret[0] == 0:
            return ret[1][3:]
        # get [X, Y, Z, Roll, Pitch, Yaw] pose
    def getpos6DoF(self):
        ret = self.robot.get_tcp_position()
        return ret[1]
    # disconnect from the robot
    def robot_disconnect(self):
        self.robot.power_off()
        print("power_off successfully")
        self.robot.logout()
        print("logout successfully")

    def jaka_connect(self):
        self.robot = jkrc.RC(self.address)
        print("logining")
        self.robot.login()
        self.robot.power_on()
        print("login successfully")
        print("power_on successfully")
# set home position
    def setHome(self, homePose = None):
        if not homePose:
            self.homePose = self.getpos6DoF()
            print(self.homePose)
        else:
            self.homePose = homePose

    def goHome(self,speedInput=5):
        ret = self.robot.linear_move(end_pos = self.homePose, move_mode = self._ABS, is_block = False, speed = speedInput)
        if(ret[0] == -4):
            print("Inverse solution failed")
        return ret
